{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b37dff",
   "metadata": {},
   "source": [
    "# NNX Demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e8099a6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "from jax import numpy as jnp\n",
    "from flax import nnx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bcc5cffe",
   "metadata": {},
   "source": [
    "### [1] NNX is Pythonic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d99b73af",
   "metadata": {
    "outputId": "d8ef66d5-6866-4d5c-94c2-d22512bfe718"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model = MLP(\n",
      "  blocks=[Block(\n",
      "      linear=Linear(\n",
      "            in_features=4,\n",
      "            out_features=4,\n",
      "            use_bias=True,\n",
      "            dtype=None,\n",
      "            param_dtype=<class 'jax.numpy.float32'>,\n",
      "            precision=None,\n",
      "            kernel_init=<function variance_scaling.<locals>.init at 0x28ae86dc0>,\n",
      "            bias_init=<function zeros at 0x122d39f70>,\n",
      "            dot_general=<function dot_general at 0x1218459d0>\n",
      "          ),\n",
      "      bn=BatchNorm(\n",
      "            num_features=4,\n",
      "  \n",
      "...\n"
     ]
    }
   ],
   "source": [
    "\n",
    "class Block(nnx.Module):\n",
    "  def __init__(self, din, dout, *, rngs):\n",
    "    self.linear = nnx.Linear(din, dout, rngs=rngs)\n",
    "    self.bn = nnx.BatchNorm(dout, rngs=rngs)\n",
    "\n",
    "  def __call__(self, x):\n",
    "    return nnx.relu(self.bn(self.linear(x)))\n",
    "\n",
    "\n",
    "class MLP(nnx.Module):\n",
    "  def __init__(self, nlayers, dim, *, rngs): # explicit RNG threading\n",
    "    self.blocks = [\n",
    "      Block(dim, dim, rngs=rngs) for _ in range(nlayers)\n",
    "    ]\n",
    "    self.count = Count(0)  # stateful variables are defined as attributes\n",
    "\n",
    "  def __call__(self, x):\n",
    "    self.count.value += 1  # in-place stateful updates\n",
    "    for block in self.blocks:\n",
    "      x = block(x)\n",
    "    return x\n",
    "\n",
    "class Count(nnx.Variable):   # custom Variable types define the \"collections\"\n",
    "  pass\n",
    "\n",
    "model = MLP(5, 4, rngs=nnx.Rngs(0))  # no special `init` method\n",
    "model.set_attributes(use_running_average=False)  # set flags\n",
    "y = model(jnp.ones((2, 4)))  # call methods directly\n",
    "\n",
    "print(f'{model = }'[:500] + '\\n...')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "523aa27c",
   "metadata": {},
   "source": [
    "Because NNX Modules contain their own state, they are very easily to inspect:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "6f278ec4",
   "metadata": {
    "outputId": "10a46b0f-2993-4677-c26d-36a4ddf33449"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "model.count = Count(\n",
      "  raw_value=1\n",
      ")\n",
      "model.blocks[0].linear.kernel = Param(\n",
      "  raw_value=Array([[-0.80345297, -0.34071913, -0.9408296 ,  0.01005968],\n",
      "         [ 0.26146442,  1.1247735 ,  0.54563737, -0.374164  ],\n",
      "         [ 1.0281805 , -0.6798804 , -0.1488401 ,  0.05694951],\n",
      "         [-0.44308168, -0.60587114,  0.434087  , -0.40541083]],      dtype=float32)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "print(f'{model.count = }')\n",
    "print(f'{model.blocks[0].linear.kernel = }')\n",
    "# print(f'{model.blocks.sdf.kernel = }') # typesafe inspection"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95f389f2",
   "metadata": {},
   "source": [
    "### [2] Model Surgery is Intuitive"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "96f61108",
   "metadata": {
    "outputId": "e6f86be8-3537-4c48-f471-316ee0fb6c45"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (2, 4)\n"
     ]
    }
   ],
   "source": [
    "# Module sharing\n",
    "model.blocks[1] = model.blocks[3]\n",
    "# Weight tying\n",
    "model.blocks[0].linear.kernel = model.blocks[-1].linear.kernel\n",
    "# Monkey patching\n",
    "def my_optimized_layer(x): return x\n",
    "model.blocks[2] = my_optimized_layer\n",
    "\n",
    "y = model(jnp.ones((2, 4)))  # still works\n",
    "print(f'{y.shape = }')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aca5a6cd",
   "metadata": {},
   "source": [
    "### [3] Interacting with JAX is easy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "c166dcc7",
   "metadata": {
    "outputId": "9a3f378b-739e-4f45-9968-574651200ede"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "state = State({\n",
      "  'blocks': {\n",
      "    '0': {\n",
      "      'linear': {\n",
      "        'kernel': Param(\n",
      "          raw_value=Array([[-0.33095378,  0.67149884,  0.33700302,  0.30972847],\n",
      "                 [ 0.8662822 , -0.11225506, -1.0820619 , -0.9906892 ],\n",
      "                 [ 0.88298297, -0.2143851 ,  0.48143268,  0.6474548 ],\n",
      "                 [-0.7710582 ,  0.3372276 ,  0.15487202,  0.6219269 ]],      dtype=float32)\n",
      "        ),\n",
      "        'bias': Param(\n",
      "          raw_value=Array([0., 0., 0., 0.], dtype=float32)\n",
      "        \n",
      "...\n",
      "\n",
      "graphdef = GraphDef(\n",
      "  type=MLP,\n",
      "  index=0,\n",
      "  attributes=('blocks', 'count'),\n",
      "  subgraphs={\n",
      "    'blocks': GraphDef(\n",
      "      type=list,\n",
      "      index=1,\n",
      "      attributes=('0', '1', '2', '3', '4'),\n",
      "      subgraphs={\n",
      "        '0': GraphDef(\n",
      "          type=Block,\n",
      "          index=2,\n",
      "          attributes=('line\n",
      "...\n"
     ]
    }
   ],
   "source": [
    "graphdef, state = model.split()\n",
    "\n",
    "# state is a dictionary-like JAX pytree\n",
    "print(f'{state = }'[:500] + '\\n...')\n",
    "\n",
    "# graphdef is also a JAX pytree, but just metadata\n",
    "print(f'\\n{graphdefefefefefef = }'[:300] + '\\n...')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "9f03e3af",
   "metadata": {
    "outputId": "0007d357-152a-449e-bcb9-b1b5a91d2d8d"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (2, 4)\n",
      "model.count.value = Array(3, dtype=int32, weak_type=True)\n"
     ]
    }
   ],
   "source": [
    "graphdef, state = model.split()\n",
    "\n",
    "@jax.jit\n",
    "def forward(graphdef: nnx.GraphDef, state: nnx.State, x: jax.Array):\n",
    "  model = graphdef.merge(state)\n",
    "  y = model(x)\n",
    "  state, _ = model.split()\n",
    "  return y, state\n",
    "\n",
    "x = jnp.ones((2, 4))\n",
    "y, state = forward(graphdef,state, x)\n",
    "\n",
    "model.update(state)\n",
    "\n",
    "print(f'{y.shape = }')\n",
    "print(f'{model.count.value = }')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "9e23dbb4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (2, 4)\n",
      "model.count = Array(4, dtype=int32, weak_type=True)\n"
     ]
    }
   ],
   "source": [
    "params, batch_stats, counts, graphdef = model.split(nnx.Param, nnx.BatchStat, Count)\n",
    "\n",
    "@jax.jit\n",
    "def forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\n",
    "  model = graphdef.merge(params, batch_stats, counts)\n",
    "  y = model(x, train=True)\n",
    "  params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\n",
    "  return y, params, batch_stats, counts\n",
    "\n",
    "x = jnp.ones((2, 4))\n",
    "y, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\n",
    "\n",
    "model.update(params, batch_stats, counts)\n",
    "\n",
    "print(f'{y.shape = }')\n",
    "print(f'{model.count = }')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "2461bfe8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "y.shape = (2, 4)\n",
      "parent.model.count.value = Array(4, dtype=int32, weak_type=True)\n"
     ]
    }
   ],
   "source": [
    "class Parent(nnx.Module):\n",
    "    def __init__(self, model: MLP):\n",
    "        self.model = model\n",
    "\n",
    "    def __call__(self, x):\n",
    "        params, batch_stats, counts, graphdef = self.model.split(nnx.Param, nnx.BatchStat, Count)\n",
    "\n",
    "        @jax.jit\n",
    "        def forward(graphdef: nnx.GraphDef, params, batch_stats, counts, x: jax.Array):\n",
    "            model = graphdef.merge(params, batch_stats, counts)\n",
    "            y = model(x)\n",
    "            params, batch_stats, counts, _ = model.split(nnx.Param, nnx.BatchStat, Count)\n",
    "            return y, params, batch_stats, counts\n",
    "\n",
    "        y, params, batch_stats, counts = forward(graphdef, params, batch_stats, counts, x)\n",
    "\n",
    "        self.model.update(params, batch_stats, counts)\n",
    "        return y\n",
    "\n",
    "parent = Parent(model)\n",
    "\n",
    "y = parent(jnp.ones((2, 4)))\n",
    "\n",
    "print(f'{y.shape = }')\n",
    "print(f'{parent.model.count.value = }')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e340bcb",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
